import os
import sys
import time

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from pyathena.async_cursor import AsyncCursor
from config import *
from pyathena import connect
import json

cursor = connect(s3_staging_dir=OutputLocation, region_name=FORMAL_S3_CONFIG['region_name'],
                 aws_access_key_id=FORMAL_S3_CONFIG['access_key'],
                 aws_secret_access_key=FORMAL_S3_CONFIG['secret_key']
                 ).cursor(schema_name=db, cursor=AsyncCursor, max_workers=2)


def submit(sql):
    result = try_submit(sql)
    if result["statusCode"] == 500:
        print("第一次sql执行失败")
        for i in range(3):
            if result["statusCode"] == 500:
                print(f"重试第{i}次执行sql")
                time.sleep(1)
                result = try_submit(sql)
            else:
                break
    if result["statusCode"] == 500:
        print("sql执行失败")
        raise Exception("执行athena sql语句异常:" + str(result["msg"]))
    else:
        return result


def try_submit(sql):
    try:
        rs = cursor.execute(sql)
        result = []
        for row in rs[1].result():
            result.append(row)
        return {
            'statusCode': 200,
            'body': json.dumps('success'),
            'msg': 'submit ok',
            'result': result
        }
    except Exception as e:
        print(f"执行sql出错:{str(e)}")
        return {
            'statusCode': 500,
            'body': json.dumps('error'),
            'msg': str(e)
        }